object_model_diagram.md:1

• The comment at line 13 appears to be template text from instructions rather than actual documentation content. This namedtuple field description doesn't match the file's purpose as an object model diagram.

# Object Model Diagram: ML Pipeline System

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

ami/ml/orchestration/nats_queue.py (1)

104-121: TTR mismatch when consumer already exists remains unaddressed.

Per the past review comment, if reserve_task() creates the consumer before publish_task(), the consumer uses the default TASK_TTR. If different TTR values are needed later, the existing consumer configuration won't be updated since JetStream doesn't allow mutating ack_wait on an existing durable consumer. This can cause premature redeliveries.

The suggested fix from the past review was to recreate the consumer when the configured ack_wait differs from the requested TTR. Consider implementing that logic or documenting that a single TTR value is always used.

ami/ml/orchestration/jobs.py (2)

44-50: Avoid duplicate url() call.

image.url() is called twice: once in the condition check and again for assignment. If this method performs any computation or I/O, this is wasteful.

     for image in images:
         image_id = str(image.pk)
-        image_url = image.url() if hasattr(image, "url") and image.url() else ""
+        image_url = image.url() if hasattr(image, "url") else ""
         if not image_url:
             job.logger.warning(f"Image {image.pk} has no URL, skipping queuing to NATS for job '{job.pk}'")
             continue

---

68-77: Inconsistent logger usage inside async context.

Lines 70 and 76 use the module-level logger, while the rest of the function uses job.logger. This inconsistency can make log correlation difficult when debugging job-specific issues.

             for image_pk, task in tasks:
                 try:
-                    logger.info(f"Queueing image {image_pk} to stream for job '{job.pk}': {task.image_url}")
+                    job.logger.info(f"Queueing image {image_pk} to stream for job '{job.pk}': {task.image_url}")
                     success = await manager.publish_task(
                         job_id=job.pk,
                         data=task,
                     )
                 except Exception as e:
-                    logger.error(f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}")
+                    job.logger.error(f"Failed to queue image {image_pk} to stream for job '{job.pk}': {e}")
                     success = False

ami/jobs/views.py (2)

299-306: Simplify redundant list comprehension.

All items in queued_tasks already have status: "queued" (set on line 289), making the filter unnecessary.

             return Response(
                 {
                     "status": "accepted",
                     "job_id": job.pk,
-                    "results_queued": len([t for t in queued_tasks if t["status"] == "queued"]),
+                    "results_queued": len(queued_tasks),
                     "tasks": queued_tasks,
                 }
             )

---

310-317: Use logger.exception() to capture traceback.

logger.exception() automatically includes the traceback, which aids debugging when errors occur in production.

         except Exception as e:
-            logger.error(f"Failed to queue pipeline results for job {job.pk}: {e}")
+            logger.exception(f"Failed to queue pipeline results for job {job.pk}: {e}")
             return Response(
                 {
                     "status": "error",
                     "job_id": job.pk,
                 },
                 status=500,
             )

ami/ml/orchestration/nats_queue.py (1)

143-144: Pydantic v1 syntax used.

.dict() is Pydantic v1 syntax. For Pydantic v2 compatibility, use .model_dump(). This applies throughout the codebase where Pydantic models are serialized.

             # Convert Pydantic model to JSON
-            task_data = json.dumps(data.dict())
+            task_data = json.dumps(data.model_dump())

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

ami/jobs/tasks.py (1)

160-160: Use timezone-aware datetime for consistency with Django settings.

The code uses datetime.now() which returns a naive datetime. The inline comment suggests this is intentional ("Use naive datetime in local time"), but Django applications with USE_TZ=True should use timezone-aware timestamps to avoid timezone-related bugs.

Apply this diff:

 def _update_job_progress(job_id: int, stage: str, progress_percentage: float) -> None:
     from ami.jobs.models import Job, JobState  # avoid circular import
+    from django.utils import timezone

     with transaction.atomic():
         job = Job.objects.select_for_update().get(pk=job_id)
         job.progress.update_stage(
             stage,
             status=JobState.SUCCESS if progress_percentage >= 1.0 else JobState.STARTED,
             progress=progress_percentage,
         )
         if stage == "results" and progress_percentage >= 1.0:
             job.status = JobState.SUCCESS
             job.progress.summary.status = JobState.SUCCESS
-            job.finished_at = datetime.now()  # Use naive datetime in local time
+            job.finished_at = timezone.now()
         job.logger.info(f"Updated job {job_id} progress in stage '{stage}' to {progress_percentage*100}%")
         job.save()

ami/jobs/tasks.py (3)

45-45: Consider narrowing the autoretry exception types.

The task uses autoretry_for=(Exception,) which will retry on any exception. This may mask distinct failure modes (e.g., validation errors vs. transient network issues) that should be handled differently. Consider specifying more targeted exceptions or using retry_for with explicit exception types.

---

123-124: Improve exception handling for NATS acknowledgment.

The code catches a bare Exception and logs via logging.error without the traceback. This makes debugging NATS acknowledgment failures difficult.

Apply this diff to use logging.exception for automatic traceback logging:

         except Exception as ack_error:
-            job.logger.error(f"Error acknowledging task via NATS: {ack_error}")
+            job.logger.exception(f"Error acknowledging task via NATS: {ack_error}")
             # Don't fail the task if ACK fails - data is already saved

Based on static analysis hints.

---

138-144: Use logging.exception for better error diagnostics.

The exception handlers at lines 139 and 142 use logging.error which doesn't include tracebacks. This makes troubleshooting production issues harder.

Apply this diff:

     except Job.DoesNotExist:
-        logger.error(f"Job {job_id} not found")
+        logger.exception(f"Job {job_id} not found")
         raise
     except Exception as e:
-        logger.error(f"Failed to process pipeline result for job {job_id}: {e}")
+        logger.exception(f"Failed to process pipeline result for job {job_id}: {e}")
         # Celery will automatically retry based on autoretry_for
         raise

Based on static analysis hints.

docker-compose.ci.yml (1)

43-55: Consider renaming container for CI consistency.

The NATS container is named ami_local_nats in a CI compose file. This could cause conflicts if both local and CI stacks run simultaneously, and doesn't follow the naming pattern of other CI resources (e.g., ami_ci_postgres_data).

Consider changing to ami_ci_nats:

   nats:
     image: nats:2.10-alpine
-    container_name: ami_local_nats
+    container_name: ami_ci_nats
     hostname: nats

Alternatively, remove the container_name directive entirely and let docker-compose auto-generate it with the stack name prefix.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

docker-compose.ci.yml (1)

19-25: Consider adding explicit health condition to depends_on for NATS.

The django service now depends on nats, but the dependency lacks condition: service_healthy, unlike the minio-init service (line 84). This means Django may start before NATS is fully initialized and healthy, potentially causing early connection failures.

If the application already has robust retry/backoff logic for NATS connections, this is acceptable. Otherwise, consider updating to ensure NATS is healthy:

     depends_on:
       - postgres
       - redis
       - minio-init
       - ml_backend
       - rabbitmq
-      - nats
+      nats:
+        condition: service_healthy

ami/ml/tests.py (3)

860-884: Solid coverage; consider relying more on the public API than _get_progress

This suite gives good coverage of initialization, per-stage tracking, empty jobs, and cleanup. The one downside is the heavy reliance on the private _get_progress helper, which couples tests to implementation details and may make refactors noisier. Where feasible, prefer going through update_state (including for test_progress_tracking and test_stages_independent) and reserve _get_progress calls for very targeted cases where you intentionally want to test that internal behavior.

---

895-918: Float equality assertions could use assertAlmostEqual for robustness

test_progress_tracking asserts exact equality on floating-point percentages (0.4, 0.8, 1.0). This works today but is a bit brittle to minor implementation changes (e.g., if total is computed differently). Using assertAlmostEqual(progress.percentage, 0.4, places=6) (and similarly for other percentages) would make these tests more resilient while preserving intent.

---

920-944: Locking behavior is covered; optional extra check for lock lifecycle

test_update_state_with_locking nicely validates that a held lock causes update_state to no-op and that progress resumes after the lock is cleared. If you want to harden this further, you could also assert that update_state does not inadvertently modify the lock when it’s held by another task (e.g., confirm the lock key/value is unchanged after the failed update) to guard against future changes to the lock-release logic.

ami/ml/orchestration/test_nats_queue.py (5)

10-38: Good structure and mocking helpers; small clarity tweaks possible

The use of IsolatedAsyncioTestCase plus _create_sample_task and _create_mock_nats_connection keeps the async tests readable and DRY. One minor improvement would be to document in _create_mock_nats_connection which subset of NATS/JetStream APIs you rely on (e.g., in a comment) so future changes to TaskQueueManager know what needs to be mocked or extended here.

---

39-49: Context manager lifecycle test is solid; optionally assert __aenter__/__aexit__ effects more strictly

The lifecycle test correctly verifies that a connection is established and closed. If you want this to catch more regressions, you could also assert that manager.nc/manager.js are None after the context exits, and optionally use nc.close.assert_awaited_once() to emphasize the async nature of the close call.

---

50-64: Publish-path test works but could assert arguments more precisely

Right now the test checks add_stream indirectly via a substring search on call_args. To make this more robust, consider asserting on the actual kwargs used (e.g., stream name and subjects) instead of relying on str(js.add_stream.call_args), which is more brittle to representation changes of mocks.

---

106-117: ACK behavior test is clear; consider also covering failure/no-connection case

The ACK test correctly ensures the manager publishes b"+ACK" to the supplied reply subject. As a follow-up, you might also add a negative test mirroring test_operations_without_connection_raise_error to ensure acknowledge_task also raises when no connection is open, keeping behavior consistent across APIs.

---

118-129: Cleanup test validates both consumer and stream deletion

test_cleanup_job_resources confirms that both delete_consumer and delete_stream are invoked and that the method returns True, which is the key behavior here. If you later add error handling (e.g., ignore “not found”), consider extending this test to cover that branch as well.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

docker-compose.staging.yml (2)

20-23: Fix indentation for the commented NATS dependency.

The commented # - nats has extra indentation compared to the other list items. While YAML allows comments at any indentation, this is inconsistent and could be confusing.

🔎 Proposed fix

     depends_on:
       - postgres
       - redis
-        # - nats  
+      # - nats

---

69-82: Consider JetStream persistence and container naming.

Two observations:

1. No JetStream persistence: Without a volume mount for JetStream storage, streams/consumers and any pending messages will be lost on container restart. This may be acceptable for staging/testing, but consider adding persistence if jobs shouldn't be lost across restarts:

   volumes:
     - nats_data:/data
   command: ["-js", "-m", "8222", "-sd", "/data"]

2. Container name inconsistency: ami_local_nats uses "local" prefix but this is the staging compose file. Consider ami_staging_nats for consistency.

Please confirm whether JetStream persistence is intentionally omitted for staging, or if it should be added to preserve task queues across container restarts.

Configuration used: defaults

Review profile: CHILL

Plan: Pro

docker-compose.staging.yml (1)

20-23: Consider uncommenting the NATS dependency or adding a note.

The nats dependency is commented out, meaning Django and Celery workers could start before NATS is ready. When the async_pipeline_workers feature flag is enabled for a project, this could cause transient connection failures during startup. If the application code includes retry logic for NATS connections, this is acceptable; otherwise, consider uncommenting the dependency.

Configuration used: defaults

Review profile: CHILL

Plan: Pro

ami/jobs/tasks.py (1)

102-130: Post-ACK failures can leave progress stale—ACK last or retry.

ACK is sent before results-stage updates; if update_state / _update_job_progress fails after ACK, NATS won’t redeliver and progress may stay behind. Consider moving ACK after results-stage updates (or explicitly retry on post-ACK failures).

🐛 Proposed fix

-        _ack_task_via_nats(reply_subject, job.logger)
-        # Update job stage with calculated progress
-        progress_info = state_manager.update_state(processed_image_ids, stage="results", request_id=self.request.id)
+        # Update job stage with calculated progress
+        progress_info = state_manager.update_state(processed_image_ids, stage="results", request_id=self.request.id)
@@
-        _update_job_progress(job_id, "results", progress_info.percentage)
+        _update_job_progress(job_id, "results", progress_info.percentage)
+        _ack_task_via_nats(reply_subject, job.logger)

Configuration used: defaults

Review profile: CHILL

Plan: Pro

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}